import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import argparse
from se_resnext import se_resnext50
# from resnext import resnext50
from data import MyDataset
from tensorboardX import SummaryWriter
import os
import sys
sys.path.append('../py2')
from params import data_path

# 定义是否使用GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 参数设置,使得我们能够手动输入命令行参数，就是让风格变得和Linux命令行差不多
parser = argparse.ArgumentParser(description='PyTorch Training')

if os.path.exists('./model') == False:
    os.makedirs('./model')

# # 输出结果保存路
parser.add_argument('--outf', default='./model', help='folder to output images and model checkpoints')
# 恢复训练时的模型路径
parser.add_argument('--net', default='./model/Resnet18.pth', help="path to net (to continue training)")

args = parser.parse_args()

# 超参数设置
EPOCH = 200  # 遍历数据集次数
pre_epoch = 0  # 定义已经遍历数据集的次数
BATCH_SIZE = 128*2  # 批处理尺寸(batch_size)
LR = 0.01  # 学习率

# 准备数据集并预处理
transform_train = transforms.Compose([
    # transforms.RandomCrop(32,padding = 4),
    transforms.RandomHorizontalFlip(),  # 图像一半的概率翻转，一半的概率不翻转
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # R,G,B每层的归一化用到的均值和方差
])

transform_test = transforms.Compose([
    # transforms.RandomCrop(32,padding = 4),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])

trainset1 = MyDataset(data_path + '/train1d.txt', transform=transform_train)
# 生成一个个batch进行批训练，组成batch的时候顺序打乱取
trainloader1 = torch.utils.data.DataLoader(trainset1, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)
testset1 = MyDataset(data_path + '/val1d.txt', transform=transform_test)
testloader1 = torch.utils.data.DataLoader(testset1, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

trainset2 = MyDataset(data_path + '/train2d.txt', transform=transform_train)
# 生成一个个batch进行批训练，组成batch的时候顺序打乱取
trainloader2 = torch.utils.data.DataLoader(trainset2, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)
testset2 = MyDataset(data_path + '/val2d.txt', transform=transform_test)
testloader2 = torch.utils.data.DataLoader(testset2, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

trainset3 = MyDataset(data_path + '/train3d.txt', transform=transform_train)
# 生成一个个batch进行批训练，组成batch的时候顺序打乱取
trainloader3 = torch.utils.data.DataLoader(trainset3, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)
testset3 = MyDataset(data_path + '/val3d.txt', transform=transform_test)
testloader3 = torch.utils.data.DataLoader(testset3, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

# 模型定义
# net = se_resnext50().to(device)
net = se_resnext50().to(device)
if torch.cuda.device_count() > 1:
    net=nn.DataParallel(net)

# 定义损失函数和优化方式
criterion = nn.CrossEntropyLoss()  # 损失函数为交叉熵，多用于多分类问题
# 优化方式为mini-batch momentum-SGD，并采用L2正则化（权重衰减）
optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9, weight_decay=5e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20100, gamma=0.8)

# 训练
if __name__ == "__main__":
    best_acc = 0  # 初始化best test accuracy
    loss_write = SummaryWriter('loss')
    acc_write = SummaryWriter('acc')
    print("Start Training!")  # 定义遍历数据集的次数
    with open("acc.txt", "w") as f1:
        with open("log.txt", "w")as f2:
            for epoch in range(pre_epoch, EPOCH):
                print('\nEpoch: %d' % (epoch + 1))
                net.train()
                sum_loss = 0.0
                correct1 = 0.0
                total1 = 0.0
                correct2 = 0.0
                total2 = 0.0
                correct3 = 0.0
                total3 = 0.0
                index = 0
                dataiter1 = iter(trainloader1)
                dataiter2 = iter(trainloader2)
                for i, data in enumerate(trainloader3, 0):
                    # 准备数据
                    length = len(trainloader3)
                    inputs, labels3 = data
                    inputs, labels3 = inputs.to(device), labels3.to(device)
                    inputs, labels2 = dataiter2.next()
                    inputs, labels1 = dataiter1.next()
                    inputs = inputs.to(device)
                    labels3 = labels3.to(device)
                    labels2 = labels2.to(device)
                    labels1 = labels1.to(device)
                    optimizer.zero_grad()
                    scheduler.step()

                    # forward + backward
                    outputs = net(inputs)

                    loss1 = criterion(outputs[0], labels1)
                    loss2 = criterion(outputs[1], labels2)
                    loss3 = criterion(outputs[2], labels3)

                    # print(labels2)

                    loss = 0.3 * loss1 + 0.3 * loss2 + 0.4 * loss3
                    # loss = loss3

                    loss.backward()
                    optimizer.step()

                    # 每训练1个batch打印一次loss和准确率
                    sum_loss += loss.item()
                    _, predicted1 = torch.max(outputs[0].data, 1)
                    total1 += labels1.size(0)
                    correct1 += predicted1.eq(labels1.data).cpu().sum()
                    Acc1 = float(10000 * correct1 / total1) / 100

                    _, predicted2 = torch.max(outputs[1].data, 1)
                    total2 += labels2.size(0)
                    correct2 += predicted2.eq(labels2.data).cpu().sum()
                    Acc2 = float(10000 * correct2 / total2) / 100

                    _, predicted3 = torch.max(outputs[2].data, 1)
                    total3 += labels3.size(0)
                    correct3 += predicted3.eq(labels3.data).cpu().sum()
                    Acc3 = float(10000 * correct3 / total3) / 100

                    print('[epoch:%d, iter:%d] Loss: %.03f | Acc1: %.2f%% | Acc2: %.2f%% | Acc3: %.2f%% '
                          % (epoch + 1, (i + 1 + epoch * length), sum_loss / (i + 1), Acc1, Acc2, Acc3))
                    f2.write('%03d  %05d |Loss: %.03f | Acc1: %.2f%% | Acc2: %.2f%% | Acc3: %.2f%% '
                             % (epoch + 1, (i + 1 + epoch * length), sum_loss / (i + 1), Acc1, Acc2, Acc3))
                    f2.write('\n')
                    f2.flush()
                    index = index + 1

                loss_write.add_scalar('loss', sum_loss / (index + 1), epoch)

                # 每训练完一个epoch测试一下准确率
                print("Waiting Test!")
                with torch.no_grad():
                    correct = 0
                    total = 0
                    for data in testloader1:
                        net.eval()
                        images, labels = data
                        images, labels = images.to(device), labels.to(device)
                        outputs = net(images)
                        # 取得分最高的那个类 (outputs.data的索引号)
                        _, predicted = torch.max(outputs[0].data, 1)
                        total += labels.size(0)
                        correct += (predicted == labels).sum()
                        Acc1_test = float(10000 * correct / total) / 100
                    print('测试分类1准确率为：%.2f%%' % Acc1_test)

                    correct = 0
                    total = 0
                    for data in testloader2:
                        net.eval()
                        images, labels = data
                        images, labels = images.to(device), labels.to(device)
                        outputs = net(images)
                        # 取得分最高的那个类 (outputs.data的索引号)
                        _, predicted = torch.max(outputs[1].data, 1)
                        total += labels.size(0)
                        correct += (predicted == labels).sum()
                        Acc2_test = float(10000 * correct / total) / 100
                    print('测试分类2准确率为：%.2f%%' % Acc2_test)

                    correct = 0
                    total = 0
                    for data in testloader3:
                        net.eval()
                        images, labels = data
                        images, labels = images.to(device), labels.to(device)
                        outputs = net(images)
                        # 取得分最高的那个类 (outputs.data的索引号)
                        _, predicted = torch.max(outputs[2].data, 1)
                        total += labels.size(0)
                        correct += (predicted == labels).sum()
                        Acc3_test = float(10000 * correct / total) / 100
                    print('测试分类3准确率为：%.2f%%' % Acc3_test)

                    acc_write.add_scalar('testacc', Acc2_test, Acc3_test, epoch)
                    # 将每次测试结果实时写入acc.txt文件中
                    print('Saving model......')
                    torch.save(net.state_dict(), '%s/net_%03d.pth' % (args.outf, epoch + 1))
                    f1.write("EPOCH=%03d,Accuracy1= %.2f%%,Accuracy2= %.2f%%,Accuracy3= %.2f%%" % (epoch + 1, Acc1_test, Acc2_test, Acc3_test))
                    f1.write('\n')
                    f1.flush()

                    # 记录最佳测试分类准确率并写入best_acc.txt文件中
                    if Acc3_test > best_acc:
                        f3 = open("best_acc.txt", "w")
                        f3.write("EPOCH=%d,best_acc1= %.2f%%, best_acc2= %.2f%%, best_acc3= %.2f%%" % (epoch + 1, Acc1_test, Acc2_test, Acc3_test))
                        f3.close()
                        best_acc = Acc3_test

            print("Training Finished, TotalEPOCH=%d" % EPOCH)
            loss_write.close()
            acc_write.close()